from coin.strategy.hk_hard_arb.products import NativeProducts


class CrossTypeArbProduct(object):
  def __init__(self, shared_product, native_products):
    self.product = shared_product
    self.native_products = native_products
    for product in native_products:
      assert shared_product.base.currency == product.base.currency, \
        (shared_product.base, product.base)


class CrossTypeArbProducts(object):
  def __init__(self,
               quote_currency_0,
               quote_currency_1,
               shared_symbols,
               symbols_0,
               symbols_1,
               ProductType_0,
               ProductType_1):
    assert len(shared_symbols) == len(symbols_0) == len(symbols_1)
    self.shared = NativeProducts.FromSharedStrProduct(quote_currency_1, shared_symbols)
    self.natives = (NativeProducts.FromTypedStrProduct(quote_currency_0, symbols_0, ProductType_0),
                    NativeProducts.FromTypedStrProduct(quote_currency_1, symbols_1, ProductType_1))
    self._arb_product_list = [
        CrossTypeArbProduct(self.shared.products[idx],
                            tuple(native.products[idx]
                                  for native in self.natives))
        for idx in range(len(shared_symbols))
    ]
    self._symbol_map = {symbol: idx for idx, symbol in enumerate(shared_symbols)}
    self._reverse_map = {}
    self._reverse_map.update({
        (0, product): idx for idx, product in enumerate(self.natives[0].products)
    })
    self._reverse_map.update({
        (1, product): idx for idx, product in enumerate(self.natives[1].products)
    })

  def by_idx(self, idx):
    return self._arb_product_list[idx]

  def by_shared_symbol(self, symbol):
    return self.by_idx(self._symbol_map[symbol])

  def by_exchange_id_and_native_product(self, exchange_id, product):
    return self.by_idx(self._reverse_map[(exchange_id, product)])
